#include <linux/module.h>
#include <linux/debugfs.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/string.h>

#define MOD_NAME "ubsan_test"
#define LOCAL_ARRSZ 20

static struct dentry *debugfs_dir;
static struct dentry *trigger_file;

enum ub_op_type {
    UB_READ,
    UB_WRITE
};

static int stack_mem_oob_right(enum ub_op_type mode)
{
    volatile char local_arr[LOCAL_ARRSZ];
    volatile char val;

    if (mode == UB_WRITE) {
        local_arr[LOCAL_ARRSZ - 5] = 'y';
        pr_info("[%s] 合法写：local_arr[%d] = 'y'\n", MOD_NAME, LOCAL_ARRSZ - 5);
        local_arr[LOCAL_ARRSZ + 5] = 'z';
        pr_info("[%s] 非法写：local_arr[%d] = 'z'（应触发UBSAN）\n", MOD_NAME, LOCAL_ARRSZ + 5);
    } else if (mode == UB_READ) {
        val = local_arr[LOCAL_ARRSZ + 8];
        pr_info("[%s] 非法读：local_arr[%d] = %c（应触发UBSAN）\n", MOD_NAME, LOCAL_ARRSZ + 8, val);
    }

    return 0;
}

static ssize_t trigger_write(struct file *file, const char __user *buf,
                             size_t count, loff_t *ppos)
{
    char cmd[32] = {0};
    size_t cmd_len = min(count, sizeof(cmd) - 1);
    char *newline = NULL;

    if (copy_from_user(cmd, buf, cmd_len))
        return -EFAULT;

    newline = strchr(cmd, '\n');
    if (newline)
        *newline = '\0';

    // 触发命令改为"1"
    if (strcmp(cmd, "1") == 0) {
        stack_mem_oob_right(UB_WRITE);
    } else {
        pr_err("[%s] 无效命令：%s（仅支持1）\n", MOD_NAME, cmd);
        return -EINVAL;
    }

    return count;
}

static const struct file_operations trigger_fops = {
    .owner = THIS_MODULE,
    .write = trigger_write,
};

static int __init ubsan_test_init(void)
{
    debugfs_dir = debugfs_create_dir(MOD_NAME, NULL);
    if (IS_ERR_OR_NULL(debugfs_dir)) {
        pr_err("[%s] debugfs目录创建失败\n", MOD_NAME);
        return PTR_ERR(debugfs_dir);
    }

    trigger_file = debugfs_create_file("trigger", 0222, debugfs_dir,
                                      NULL, &trigger_fops);
    if (IS_ERR_OR_NULL(trigger_file)) {
        pr_err("[%s] trigger文件创建失败\n", MOD_NAME);
        debugfs_remove(debugfs_dir);
        return PTR_ERR(trigger_file);
    }

    pr_info("[%s] 模块加载成功！触发接口：/sys/kernel/debug/%s/trigger\n",
            MOD_NAME, MOD_NAME);
    pr_info("[%s] 支持命令：echo 1 > trigger（栈数组越界写）\n", MOD_NAME);
    return 0;
}

static void __exit ubsan_test_exit(void)
{
    debugfs_remove(trigger_file);
    debugfs_remove(debugfs_dir);
    pr_info("[%s] 模块卸载成功\n", MOD_NAME);
}

module_init(ubsan_test_init);
module_exit(ubsan_test_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("UBSAN Test Module: Catch Stack OOB UB");
MODULE_AUTHOR("Your Name");

